import matplotlib.pyplot as plt
import io
import tqdm
import random
import numpy as np
from PIL import Image


class CellMachine:
    def __init__(self, width: int, height: int):
        self.width  = width
        self.height = height
        self.cells  = [[False] * width for _ in range(height)]

    def randomize(self, p_true = 0.5):
        for x in range(self.height):
            for y in range(self.width):
                self.cells[x][y] = np.random.choice([True, False], p=[p_true, 1 - p_true])

    def next_epoch(self, rule: dict = None):
        rule = rule if rule else self.get_default_rule()
        assert all([x not in rule['D'] for x in rule['S']])
        next_cells = [[False] * self.width for _ in range(self.height)]
        for x in range(self.height):
            for y in range(self.width):
                num_of_neighbors = self.get_num_of_neighbors(x, y, rule)

                if  self.cells[x][y]:
                    if num_of_neighbors in rule['S']:
                        next_cells[x][y] = True
                    elif num_of_neighbors in rule['D']:
                        next_cells[x][y] = False
                elif num_of_neighbors in rule['B']:
                    next_cells[x][y] = True
        self.cells = next_cells

    def get_num_of_neighbors(self, center_x: int, center_y: int, rule: dict) -> int:
        result = 0
        for raw_x in range(center_x - rule['R'], center_x + rule['R'] + 1):
            for raw_y in range(center_y - rule['R'], center_y + rule['R'] + 1):
                if   rule['N'] == 'Von Neumann':
                    if abs(raw_x - center_x) + abs(raw_y - center_y) <= rule['R']:
                        result += int(self.cells[raw_x % self.height][raw_y % self.width])

                elif rule['N'] == 'Moore':
                    result += int(self.cells[raw_x % self.height][raw_y % self.width])

                else:
                    raise ValueError('Invalid neighborhood type')
        return result

    def pretty(self, true='.', false=' '):
        result = ''
        for row in self.cells:
            for value in row:
                result += f'{true if value else false} '
            result += '\n'
        return result

    @staticmethod
    def get_default_rule():
        return {
                'S': [1, 2, 3],
                'D': [4, 5, 6, 7, 8, 9],
                'B': [1],
                'R': 1,
                'N': 'Moore'  # Moore, Von Neumann
            }

    def make_gif(self, n = 10000, rule: dict = None, figsize: tuple[float, float] = (10, 10), s: int = 10):
        gif = []
        rule = rule if rule else self.get_default_rule()
        for epoch in tqdm.tqdm(range(n), desc='Making Gif'):
            self.next_epoch(rule)
            fig, ax = plt.subplots(figsize=figsize)

            black_x, black_y = [], []
            white_x, white_y = [], []

            for x in range(self.height):
                for y in range(self.width):
                    if self.cells[x][y]:
                        black_x.append(x)
                        black_y.append(y)
                    else:
                        white_x.append(x)
                        white_y.append(y)

            ax.scatter(black_x, black_y, color='black', s=s, marker='s')
            # ax.scatter(white_x, white_y, color='white', s=s, marker='s')
            ax.axis('off')
            ax.set_title(f'Epoch {epoch}')
            gif.append(self.fig2img(fig))
            plt.close(fig)

        gif[0].save(
            'cell_machine.gif',
            save_all      = True,
            append_images = gif[1:],
            optimize      = True,
            duration      = 500,
            loop          = 0
        )

    @staticmethod
    def fig2img(fig):

        buf = io.BytesIO()
        fig.savefig(buf)
        buf.seek(0)
        img = Image.open(buf)
        return img